// SPDX-License-Identifier: Mulan PSL v2
/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This software is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *         http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

use std::sync::Arc;
use crate::mock_common::ipc::message::{Request, Argument};
use crate::error::{ServerError, ServerResult};
use crate::gpu::{GPUResourceManager, DevicePtr};

/// CUDA API method identifiers.
/// These should match the method IDs used by the proxy.
pub mod method_ids {
    pub const CUDA_MALLOC: u64 = 1;
    pub const CUDA_FREE: u64 = 2;
    pub const CUDA_MEMCPY: u64 = 3;
    pub const CUDA_MEMCPY_H2D: u64 = 4;
    pub const CUDA_MEMCPY_D2H: u64 = 5;
    pub const CUDA_GET_DEVICE_COUNT: u64 = 6;
    pub const CUDA_GET_DEVICE_INFO: u64 = 7;
    pub const CUDA_SET_DEVICE: u64 = 8;
    pub const CUDA_GET_DEVICE: u64 = 9;
    pub const CUDA_DEVICE_SYNCHRONIZE: u64 = 10;
}

/// Response for CUDA API calls.
#[derive(Debug, Clone)]
pub struct ApiResponse {
    pub success: bool,
    pub error_code: i32,
    pub error_message: Option<String>,
    pub data: Vec<u8>,
}

impl ApiResponse {
    /// Create a successful response.
    pub fn success(data: Vec<u8>) -> Self {
        Self {
            success: true,
            error_code: 0,
            error_message: None,
            data,
        }
    }
    
    /// Create an error response.
    pub fn error(code: i32, message: String) -> Self {
        Self {
            success: false,
            error_code: code,
            error_message: Some(message),
            data: Vec::new(),
        }
    }
    
    /// Create a response from a ServerError.
    pub fn from_error(error: ServerError) -> Self {
        Self::error(-1, error.to_string())
    }
}

/// CUDA API handler responsible for processing CUDA API requests.
#[derive(Clone)]
pub struct CudaApiHandler {
    gpu_manager: Arc<GPUResourceManager>,
    current_device: std::sync::Arc<std::sync::RwLock<u32>>,
}

impl CudaApiHandler {
    /// Create a new CUDA API handler.
    pub fn new(gpu_manager: Arc<GPUResourceManager>) -> Self {
        Self {
            gpu_manager,
            current_device: std::sync::Arc::new(std::sync::RwLock::new(0)),
        }
    }
    
    /// Handle a CUDA API request.
    pub fn handle_request(&self, request: &Request) -> ServerResult<ApiResponse> {
        match request.method_id() {
            method_ids::CUDA_MALLOC => self.handle_cuda_malloc(request),
            method_ids::CUDA_FREE => self.handle_cuda_free(request),
            method_ids::CUDA_MEMCPY => self.handle_cuda_memcpy(request),
            method_ids::CUDA_MEMCPY_H2D => self.handle_cuda_memcpy_h2d(request),
            method_ids::CUDA_MEMCPY_D2H => self.handle_cuda_memcpy_d2h(request),
            method_ids::CUDA_GET_DEVICE_COUNT => self.handle_get_device_count(request),
            method_ids::CUDA_GET_DEVICE_INFO => self.handle_get_device_info(request),
            method_ids::CUDA_SET_DEVICE => self.handle_set_device(request),
            method_ids::CUDA_GET_DEVICE => self.handle_get_device(request),
            method_ids::CUDA_DEVICE_SYNCHRONIZE => self.handle_device_synchronize(request),
            _ => Err(ServerError::request_processing(format!(
                "Unknown method ID: {}", request.method_id()
            ))),
        }
    }
    
    /// Handle cudaMalloc request.
    fn handle_cuda_malloc(&self, request: &Request) -> ServerResult<ApiResponse> {
        if request.argc() != 1 {
            return Ok(ApiResponse::error(
                -1, 
                format!("cudaMalloc expects 1 argument, got {}", request.argc())
            ));
        }
        
        let size_arg = request.argv(0)
            .ok_or_else(|| ServerError::request_processing("Missing size argument".to_string()))?;
        
        let size = self.extract_size_from_arg(size_arg)?;
        let device_id = *self.current_device.read().unwrap();
        
        match self.gpu_manager.allocate_memory(size, device_id) {
            Ok(ptr) => {
                log::debug!("cudaMalloc: allocated {} bytes at 0x{:x} on device {}", size, ptr, device_id);
                Ok(ApiResponse::success(ptr.to_le_bytes().to_vec()))
            },
            Err(e) => Ok(ApiResponse::from_error(e)),
        }
    }
    
    /// Handle cudaFree request.
    fn handle_cuda_free(&self, request: &Request) -> ServerResult<ApiResponse> {
        if request.argc() != 1 {
            return Ok(ApiResponse::error(
                -1,
                format!("cudaFree expects 1 argument, got {}", request.argc())
            ));
        }
        
        let ptr_arg = request.argv(0)
            .ok_or_else(|| ServerError::request_processing("Missing pointer argument".to_string()))?;
        
        let ptr = self.extract_ptr_from_arg(ptr_arg)?;
        
        match self.gpu_manager.free_memory(ptr) {
            Ok(()) => {
                log::debug!("cudaFree: freed memory at 0x{:x}", ptr);
                Ok(ApiResponse::success(vec![0]))
            },
            Err(e) => Ok(ApiResponse::from_error(e)),
        }
    }
    
    /// Handle cudaMemcpy request.
    fn handle_cuda_memcpy(&self, request: &Request) -> ServerResult<ApiResponse> {
        if request.argc() != 3 {
            return Ok(ApiResponse::error(
                -1,
                format!("cudaMemcpy expects 3 arguments, got {}", request.argc())
            ));
        }
        
        let dst_arg = request.argv(0).unwrap();
        let src_arg = request.argv(1).unwrap();
        let size_arg = request.argv(2).unwrap();
        
        let dst = self.extract_ptr_from_arg(dst_arg)?;
        let src = self.extract_ptr_from_arg(src_arg)?;
        let size = self.extract_size_from_arg(size_arg)?;
        
        match self.gpu_manager.copy_memory(src, dst, size) {
            Ok(()) => {
                log::debug!("cudaMemcpy: copied {} bytes from 0x{:x} to 0x{:x}", size, src, dst);
                Ok(ApiResponse::success(vec![0]))
            },
            Err(e) => Ok(ApiResponse::from_error(e)),
        }
    }
    
    /// Handle cudaMemcpy host-to-device request.
    fn handle_cuda_memcpy_h2d(&self, request: &Request) -> ServerResult<ApiResponse> {
        if request.argc() != 3 {
            return Ok(ApiResponse::error(
                -1,
                format!("cudaMemcpyH2D expects 3 arguments, got {}", request.argc())
            ));
        }
        
        let dst_arg = request.argv(0).unwrap();
        let src_data_arg = request.argv(1).unwrap();
        let size_arg = request.argv(2).unwrap();
        
        let _dst = self.extract_ptr_from_arg(dst_arg)?;
        let _src_data = self.extract_data_from_arg(src_data_arg)?;
        let _size = self.extract_size_from_arg(size_arg)?;
        
        // Mock implementation - in real version would copy data to GPU
        log::debug!("cudaMemcpyH2D: mock host-to-device copy");
        Ok(ApiResponse::success(vec![0]))
    }
    
    /// Handle cudaMemcpy device-to-host request.
    fn handle_cuda_memcpy_d2h(&self, request: &Request) -> ServerResult<ApiResponse> {
        if request.argc() != 3 {
            return Ok(ApiResponse::error(
                -1,
                format!("cudaMemcpyD2H expects 3 arguments, got {}", request.argc())
            ));
        }
        
        let _dst_data_arg = request.argv(0).unwrap();
        let _src_arg = request.argv(1).unwrap();
        let size_arg = request.argv(2).unwrap();
        
        let size = self.extract_size_from_arg(size_arg)?;
        
        // Mock implementation - in real version would copy data from GPU
        log::debug!("cudaMemcpyD2H: mock device-to-host copy of {} bytes", size);
        
        // Return mock data
        let mock_data = vec![0u8; size];
        Ok(ApiResponse::success(mock_data))
    }
    
    /// Handle cudaGetDeviceCount request.
    fn handle_get_device_count(&self, _request: &Request) -> ServerResult<ApiResponse> {
        let devices = self.gpu_manager.get_available_devices();
        let count = devices.len() as u32;
        
        log::debug!("cudaGetDeviceCount: {} devices", count);
        Ok(ApiResponse::success(count.to_le_bytes().to_vec()))
    }
    
    /// Handle cudaGetDeviceProperties request.
    fn handle_get_device_info(&self, request: &Request) -> ServerResult<ApiResponse> {
        if request.argc() != 1 {
            return Ok(ApiResponse::error(
                -1,
                format!("cudaGetDeviceProperties expects 1 argument, got {}", request.argc())
            ));
        }
        
        let device_arg = request.argv(0).unwrap();
        let device_id = self.extract_u32_from_arg(device_arg)?;
        
        match self.gpu_manager.get_device_info(device_id) {
            Ok(info) => {
                log::debug!("cudaGetDeviceProperties: device {} info", device_id);
                
                // Serialize device info (simplified)
                let mut data = Vec::new();
                data.extend_from_slice(&info.device_id.to_le_bytes());
                data.extend_from_slice(&info.total_memory.to_le_bytes());
                data.extend_from_slice(&info.free_memory.to_le_bytes());
                data.extend_from_slice(&info.compute_capability_major.to_le_bytes());
                data.extend_from_slice(&info.compute_capability_minor.to_le_bytes());
                data.extend_from_slice(&info.max_threads_per_block.to_le_bytes());
                data.extend_from_slice(&info.max_blocks_per_grid.to_le_bytes());
                
                // Add device name
                let name_bytes = info.name.as_bytes();
                data.extend_from_slice(&(name_bytes.len() as u32).to_le_bytes());
                data.extend_from_slice(name_bytes);
                
                Ok(ApiResponse::success(data))
            },
            Err(e) => Ok(ApiResponse::from_error(e)),
        }
    }
    
    /// Handle cudaSetDevice request.
    fn handle_set_device(&self, request: &Request) -> ServerResult<ApiResponse> {
        if request.argc() != 1 {
            return Ok(ApiResponse::error(
                -1,
                format!("cudaSetDevice expects 1 argument, got {}", request.argc())
            ));
        }
        
        let device_arg = request.argv(0).unwrap();
        let device_id = self.extract_u32_from_arg(device_arg)?;
        
        // Verify device exists
        if self.gpu_manager.get_device_info(device_id).is_err() {
            return Ok(ApiResponse::error(
                -1,
                format!("Device {} not available", device_id)
            ));
        }
        
        *self.current_device.write().unwrap() = device_id;
        log::debug!("cudaSetDevice: set current device to {}", device_id);
        
        Ok(ApiResponse::success(vec![0]))
    }
    
    /// Handle cudaGetDevice request.
    fn handle_get_device(&self, _request: &Request) -> ServerResult<ApiResponse> {
        let device_id = *self.current_device.read().unwrap();
        log::debug!("cudaGetDevice: current device is {}", device_id);
        
        Ok(ApiResponse::success(device_id.to_le_bytes().to_vec()))
    }
    
    /// Handle cudaDeviceSynchronize request.
    fn handle_device_synchronize(&self, _request: &Request) -> ServerResult<ApiResponse> {
        let device_id = *self.current_device.read().unwrap();
        log::debug!("cudaDeviceSynchronize: synchronizing device {}", device_id);
        
        // Mock implementation - in real version would call cudaDeviceSynchronize
        Ok(ApiResponse::success(vec![0]))
    }
    
    /// Extract size from argument.
    fn extract_size_from_arg(&self, arg: &Argument) -> ServerResult<usize> {
        // This is a simplified extraction - in real implementation,
        // we would need to properly handle the Argument type
        // For now, assume the argument contains the size as bytes
        if let Some(data) = self.get_arg_data(arg) {
            if data.len() >= 8 {
                let size_bytes: [u8; 8] = data[0..8].try_into().unwrap();
                return Ok(usize::from_le_bytes(size_bytes));
            }
        }
        Err(ServerError::request_processing("Invalid size argument".to_string()))
    }
    
    /// Extract device pointer from argument.
    fn extract_ptr_from_arg(&self, arg: &Argument) -> ServerResult<DevicePtr> {
        if let Some(data) = self.get_arg_data(arg) {
            if data.len() >= 8 {
                let ptr_bytes: [u8; 8] = data[0..8].try_into().unwrap();
                return Ok(DevicePtr::from_le_bytes(ptr_bytes));
            }
        }
        Err(ServerError::request_processing("Invalid pointer argument".to_string()))
    }
    
    /// Extract u32 from argument.
    fn extract_u32_from_arg(&self, arg: &Argument) -> ServerResult<u32> {
        if let Some(data) = self.get_arg_data(arg) {
            if data.len() >= 4 {
                let value_bytes: [u8; 4] = data[0..4].try_into().unwrap();
                return Ok(u32::from_le_bytes(value_bytes));
            }
        }
        Err(ServerError::request_processing("Invalid u32 argument".to_string()))
    }
    
    /// Extract data from argument.
    fn extract_data_from_arg(&self, arg: &Argument) -> ServerResult<Vec<u8>> {
        if let Some(data) = self.get_arg_data(arg) {
            return Ok(data.to_vec());
        }
        Err(ServerError::request_processing("Invalid data argument".to_string()))
    }
    
    /// Get data from argument (placeholder implementation).
    fn get_arg_data(&self, _arg: &Argument) -> Option<&[u8]> {
        // This is a placeholder - in real implementation we would
        // extract data from the Argument based on its actual structure
        None
    }
}